package bloom_filter

import (
    "encoding/gob"
    "github.com/dgryski/go-farm"
    "math/rand"
    "os"
    "bufio"
)

type BloomFilter struct {
    Arr      []byte
    BitCount uint
    Seeds    []uint32
}

func NewBloomFilter(hashNum, arrayLen int) *BloomFilter {
    size := arrayLen / 8
    seeds := make([]uint32, hashNum) //随机生成5个种子
    for i := 0; i < hashNum; i++ {
        seeds[i] = uint32(rand.Intn(1000000000))
    }
    return &BloomFilter{
        Arr:      make([]byte, size),
        BitCount: uint(size * 8),
        Seeds:    seeds,
    }
}

// 判断bit数组的index位置是否为1
func (bf *BloomFilter) getBit(index uint) bool {
    var a uint = index / 8
    var bt uint = uint(bf.Arr[a])

    var b uint = index % 8
    var c uint = 1 << b
    return bt&c == c
}

// 设置bit数组的index位置为1
func (bf *BloomFilter) setBit(index uint) {
    var a uint = index / 8
    var bt uint = uint(bf.Arr[a])

    var b uint = index % 8
    var c uint = 1 << b
    bf.Arr[a] = byte(bt | c)
}

// 添加一个元素到BloomFilter中
func (bf *BloomFilter) Add(value string) {
    for range bf.Seeds {
        hashValue := farm.Fingerprint32([]byte(value))
        index := hashValue % uint32(bf.BitCount)
        bf.setBit(uint(index))
    }
}

// 判断一个元素是否存在BloomFilter中
func (bf *BloomFilter) Exists(value string) bool {
    for range bf.Seeds {
        hashValue := farm.Fingerprint32([]byte(value))
        index := hashValue % uint32(bf.BitCount)
        if !bf.getBit(uint(index)) {
            return false
        }
    }
    return true
}

// 将BloomFilter导出到文件中
func (bf *BloomFilter) Dump(outfile string) error {
    fout, err := os.OpenFile(outfile, os.O_CREATE|os.O_WRONLY, 0666)
    if err != nil {
        return err
    }
    defer fout.Close()
    writer := bufio.NewWriter(fout)
    defer writer.Flush()
    encoder := gob.NewEncoder(writer)
    return encoder.Encode(bf)
}

// 从文件中加载BloomFilter
func LoadBloomFilter(infile string) (*BloomFilter, error) {
    fin, err := os.Open(infile)
    if err != nil {
        return nil, err
    }
    defer fin.Close()
    reader := bufio.NewReader(fin)
    decoder := gob.NewDecoder(reader)
    var bf BloomFilter
    err = decoder.Decode(&bf)
    if err != nil {
        return nil, err
    }
    return &bf, nil
}